# -*- coding=utf-8 -*-

import pandas as pd
from sqlmodel import SQLModel, Field, Session, create_engine
from pathlib import Path
from typing import Optional, List, Dict, Any
import warnings

# 禁用警告
warnings.filterwarnings('ignore', category=UserWarning)
warnings.simplefilter(action='ignore', category=pd.errors.SettingWithCopyWarning)

# 定义数据库模型
class Record(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    类型: str
    顾客代码: Optional[str] = Field(default=None, alias="顾客代码")
    下单时间: Optional[str] = None
    销售回复出货时间: Optional[str] = None
    实际出货时间: Optional[str] = None
    采购单号: Optional[str] = None
    中聚订单号: Optional[str] = None
    电池型号: Optional[str] = None
    容量mash: Optional[float] = None
    技术图纸: Optional[str] = None
    带电量: Optional[float] = None
    下单数量: Optional[int] = None
    实发数量: Optional[int] = None
    入库数量: Optional[int] = None
    单价: Optional[float] = None
    总价: Optional[float] = None
    备注: Optional[str] = None

def get_excel_files(fpath: str) -> List[Path]:
    """获取有效Excel文件列表"""
    path = Path(fpath)
    
    if path.is_file():
        return [path] if not path.name.startswith("~$") else []
    elif path.is_dir():
        valid_files = []
        for ext in ["*.xlsx", "*.xls"]:
            for f in path.glob(ext):
                # 过滤临时文件和隐藏文件
                if not f.name.startswith(("~$", ".")):
                    valid_files.append(f)
        return valid_files
    else:
        raise ValueError(f"无效路径: {fpath}")

def clean_column_name(col: Any) -> str:
    """统一清理列名"""
    return str(col).strip().replace("\n", "").replace(" ", "").lower()

def map_to_model_fields(data_df: pd.DataFrame) -> pd.DataFrame:
    """将DataFrame列映射到模型字段"""
    # 创建字段映射字典（小写键）
    field_mapping = {
        "顾客代码": "顾客代码",
        "下单时间": "下单时间",
        "销售回复出货时间": "销售回复出货时间",
        "实际出货时间": "实际出货时间",
        "采购单号": "采购单号",
        "中聚订单号": "中聚订单号",
        "电池型号": "电池型号",
        "容量mash": "容量mash",
        "技术图纸": "技术图纸",
        "带电量": "带电量",
        "下单数量": "下单数量",
        "实发数量": "实发数量",
        "入库数量": "入库数量",
        "单价": "单价",
        "总价": "总价",
        "备注": "备注", 
    }
    lower_mapping = {k.lower(): v for k, v in field_mapping.items()}
    
    # 创建数据副本避免警告
    df = data_df.copy()
    
    # 清理列名并映射
    new_columns = []
    for col in df.columns:
        cleaned = clean_column_name(col)
        new_columns.append(lower_mapping.get(cleaned, None))
    df.columns = new_columns
    
    # 添加缺失字段
    model_fields = [f for f in Record.__annotations__ if f not in ("id", "类型")]
    for field in model_fields:
        if field not in df.columns:
            df = df.assign(**{field: None})  # 安全添加列
    
    return df[model_fields].copy()

def process_sheet(file_path: Path, sheet_name: str, engine) -> None:
    """处理单个工作表"""
    try:
        # 读取原始数据
        raw_df = pd.read_excel(
            file_path,
            sheet_name=sheet_name,
            header=None,
            dtype=str,
            keep_default_na=False
        )
        
        # 空表检测
        if raw_df.empty or len(raw_df) < 3:
            print(f"↳ 工作表 {sheet_name} 格式无效，已跳过")
            return
        
        # 获取表头（第三行）
        header_row = raw_df.iloc[2].apply(clean_column_name).tolist()
        
        # 检查有效表头
        if not any(header_row):
            print(f"↳ 工作表 {sheet_name} 缺少有效表头，已跳过")
            return
        
        # 定位数据起始和结束位置
        data_start = 3
        end_mask = raw_df.iloc[data_start:, 0].apply(
            lambda x: str(x).strip() == ""
        )
        data_end = end_mask.idxmax() if end_mask.any() else len(raw_df)
        
        # 提取有效数据并创建副本
        data_df = raw_df.iloc[data_start:data_end].copy()
        data_df.columns = header_row
        
        # 字段映射
        processed_df = map_to_model_fields(data_df)
        
        # 移除可能的类型字段
        if "类型" in processed_df.columns:
            processed_df = processed_df.drop(columns=["类型"])
        
        # 添加类型字段
        sheet_type = sheet_name[-2:] if len(sheet_name) >= 2 else sheet_name
        processed_df.insert(0, "类型", sheet_type)
        
        # 类型转换
        numeric_cols = ["容量mash", "带电量", "下单数量", "实发数量", 
                       "入库数量", "单价", "总价"]
        for col in numeric_cols:
            if col in processed_df.columns:
                processed_df[col] = pd.to_numeric(
                    processed_df[col], 
                    errors="coerce"
                ).fillna(0)
        
        # 写入数据库
        with Session(engine) as session:
            # 分批次提交（每200条）
            for i in range(0, len(processed_df), 200):
                batch = processed_df.iloc[i:i+200]
                records = []
                for _, row in batch.iterrows():
                    try:
                        record_data = {
                            k: v if pd.notnull(v) else None
                            for k, v in row.items()
                        }
                        records.append(Record(**record_data))
                    except Exception as e:
                        print(f"↳ 记录转换失败: {str(e)}")
                        continue
                if records:
                    session.add_all(records)
                    session.commit()
        
        print(f"↳ 成功导入 {len(processed_df)} 条记录")
        
    except Exception as e:
        print(f"↳ 处理工作表 {sheet_name} 时出错: {str(e)}")

def main(fpath: str) -> None:
    """主处理函数"""
    engine = create_engine("sqlite:///msbase.db")
    SQLModel.metadata.create_all(engine)
    
    try:
        files = get_excel_files(fpath)
        print(f"发现 {len(files)} 个有效文件")
    except Exception as e:
        print(f"路径错误: {str(e)}")
        return
    
    for file_path in files:
        try:
            print(f"\n正在处理文件: {file_path.name}")
            with pd.ExcelFile(file_path) as xls:
                sheets = xls.sheet_names
                print(f"包含 {len(sheets)} 个工作表")
                for i, sheet_name in enumerate(sheets, 1):
                    print(f"├─ 处理工作表中 ({i}/{len(sheets)})：{sheet_name}")
                    process_sheet(file_path, sheet_name, engine)
        except Exception as e:
            print(f"└─ 文件处理失败: {str(e)}")
    
    print("\n✅ 数据处理完成")

if __name__ == "__main__":
    # 请在此处指定要处理的文件夹路径
    folder_path = r"E:\Documents\Desktop\two2025\2025年对账单"
    main(folder_path)
